"""
A simple DICOM reader for 16 bit gray-scale images.

  Support for sequence and encapsulation has not been
  implemented yet.

  Only tested for MR and CT images.

2001-05-04
Hua Qian
hqian@irus.rri.on.ca

"""

import struct
import numarray
import string
import sys
# import sys

from dicomdict import dicomdict    

class dicom:
    def __init__(self,filename=None):
        if filename:
            self.setFileName(filename)

    def setFileName(self, filename):
        self._FileName = filename

        # DICOM header info
        self._info = {}
        self._infoShort = {}

        # Fill self._info and self._infoShort with necessary default values
        self._info['RowSpacing'] = 1
        self._info['ColumnSpacing'] = 1
        self._info['RowOrientation'] = (1.0, 0.0, 0.0)
        self._info['ColumnOrientation'] = (0.0, 1.0, 0.0)
        self._info['Position'] = (0.0, 0.0, 0.0)
        self._infoShort['Position'] = (0.0, 0.0, 0.0)
        
        # pixel data 
        self.PixelDataEndian = 'little'
        self.PixelDataType = 's'   # numarray Array type code
        self.PixelDataOffset = None
        self.PixelDataSwapFlag = 0
        self._PixelData = None

    def dicomTags(self, loaddata=0):
        """Return a dictionary with all recognized DICOM tags,
        except that PixelData will not be present but instead
        a PixelDataOffset tag will specify the offset in the
        file where the data starts.
        """
        self._LoadPixelDataFlag = loaddata

        self._ReadDICOMTags()
        return self._tags
             
    def dicomInfo(self, loaddata=0):
        """ return a dictionary with patient, study, series
        and image information, such as 'PatientName', 'Rows',
        'Columns', etc.
        """
        
        self._LoadPixelDataFlag = loaddata
        
        self._ReadDICOMInfo()
        if( self._File == None ):
        	return {}
        else:
        	return self._info

    def dicomInfoShort(self, loaddata=0):
        """ return a short dictionary with following keys
        'Position', 'RowOrientation', 'ColumnOrientation',
        'ImageType', 'PixelDataOffset'.
        """       
        self._LoadPixelDataFlag = loaddata

        self._ReadDICOMInfoShort()
        return self._infoShort

    def getImage(self):
        """Return an array (numarray module) of the pixel data
        that is properly swapped for byte order.
        """
        
        if self._PixelData:
            print "pixel data already loaded"
            return self._PixelData
        else:
            #self.dicomInfo(loaddata=1)
            print "reading data"
            self.dicomTags(loaddata=1)
            return self._PixelData

    def _DICOMOpen(self):
        """ open DICOM file and read some meta information
        so that we are ready to read data elements group-tag-value.
        """
        
        # some variables for DICOM Data Element
        self._Group = 0
        self._Element = 0
        self._VR = None         # DICOM Value Representation
        self._Length = 0        # number of bytes of the value
        self._Value = None
        self._VM = 1            # DICOM Value Multiplility
        self._Encapsulated = 0
        self._Sequence = 0

        self._DICOMTransferSyntax = None
        self._MetaSyntax = 'implicit'
        self._DataSyntax = 'implicit'
        self._MetaEndian = 'little'
        self._DataEndian = 'little'
        
        self._EndianChar = '<'  # Python struct pack/unpack format

        self._Meta = 0          # Meta group flag 

        # open the file and browse the meta info
        self._File = open(self._FileName, 'rb')
        self._File.seek(128,0)
        b = self._File.read(4)
        if b == 'DICM':
            self._Meta = 1
            b = self._File.read(6)
            if (string.find(string.uppercase, b[4]) != -1 and
                string.find(string.uppercase, b[5]) != -1):
                self._MetaSyntax = 'explicit'
                self._DataSyntax = 'explicit'

            # rewind file position to the meta information area 
            self._File.seek(132,0)
                
            # print "Part 10 file: ", self._MetaSyntax, self._DataSyntax, self._EndianChar

        else:  
	        print "Not a DICOM file"
        	self._File.close()
        	self._File = None
        	return # not stardard DICOM part 10 file, no meta info
        	"""self._File.seek(0,0)
            b = self._File.read(8)
            # examine probable group number ... assume <= 0x00FF
            if (b[0] < b[1]):
                self._DataEndian = 'big'
                # self._EndianChar = '>'
                # self.PixelDataEndian = 'big'
            else:
                # group number is zero
                # try the 32 bit value length of implicit VR
                if (b[0] == 0 and b[1] ==0 and b[4] < b[7]):
                    self._DataEndian = 'big'
                    # self._EndianChar = '>'
                    # self.PixelDataEndian = 'big'
            if (string.find(string.uppercase, b[4]) == 1 and
                string.find(string.uppercase, b[5]) == 1):
                self._DataSyntax = 'explicit'
            # else unrecognized ... assume default 'implicit'

            # rewind file position to the beginning
            self._File.seek(0,0)

            # print "None DICOM part 10 file:", self._MetaSyntax, self._DataSyntax, self._EndianChar"""
  
    def _DICOMClose(self):
        self._Value = None
        self._File.close()

    def _DICOMReadTag(self):
        # print self._EndianChar
        if self._Meta:
            syntax = self._MetaSyntax
            endian = self._MetaEndian
        else:
            syntax = self._DataSyntax
            endian = self._DataEndian
        
        if endian == 'big':
            self._EndianChar = '>'
        else:
            self._EndianChar = '<'
            
        buf = self._File.read(4)
        if len(buf) == 0:
            return 0
        
        self._Group, self._Element = struct.unpack(self._EndianChar + 'HH', buf)        
        # check if we are out of meta info section
        if self._Meta:
            syntax = self._MetaSyntax
            endian = self._MetaEndian
            if self._Group != 0x0002:
                self._Meta = 0
                syntax = self._DataSyntax
                endian = self._DataEndian
                 
                # uppack the value again using use endian
                if endian == 'big':
                    self._EndianChar = '>'
                else:
                    self._EndianChar = '<'                                           
                self._Group, self._Element = struct.unpack(self._EndianChar + 'HH', buf)
        else :
            syntax = self._DataSyntax
            endian = self._DataEndian

        if endian == 'big':
            self._EndianChar = '>'
        else:
            self._EndianChar = '<'                                  
        
        if(syntax == 'implicit' or self._Group == 0xFFFE):
            self._VR = None
            """
            Here, we should find the VR from dictionary
               self._GetVR()
            The dictionary will be implemented later.
            For now, we input the VR when we read the tag value
            """ 
            
            buf = self._File.read(4)
            self._Length = struct.unpack(self._EndianChar+ 'I', buf)[0]

        else:
            vr = self._File.read(2)
            if (vr == 'OB' or
                vr == 'OW' or
                vr == 'SQ' or
                vr == 'UN' or
                vr == 'UT'):
                self._File.seek(2,1)
                buf = self._File.read(4)
                self._Length = struct.unpack(self._EndianChar + 'I', buf)[0]
            else:
                buf = self._File.read(2)
                self._Length = struct.unpack(self._EndianChar + 'H', buf)[0]

            self._VR = vr

        # find out data transfer syntax
        if self._Group == 0x0002 and self._Element == 0x0010:
            saved_fp = self._File.tell()
            self._DICOMReadValue('UI')

            self._DICOMDataTransferSyntax = self._Value
            # rewind
            self._File.seek(saved_fp, 0)

            if string.find(self._Value, '1.2.840.113619.5.2') == 0 :
                # GE Implicity VR Little Endian except Big endian pixels
                self._DataSyntax = 'implicit'
                self._DataEndian = 'little'
                self.PixelDataEndian = 'big'
            elif string.find(self._Value, '1.2.840.10008.1.2.1') == 0 :
                # Explicit VR Little Endian
                self._DataSyntax = 'explicit'
                self._DataEndian = 'little'
                self.PixelDataEndian = 'little'
            elif string.find(self._Value, '1.2.840.10008.1.2.2') == 0 :
                # Explicit VR Big Endian
                self._DataSyntax = 'explicit'
                self._DataEndian = 'big'
                self.PixelDataEndian = 'big'
            elif string.find(self._Value, '1.2.840.10008.1.2.4') == 0 :
                print "warning: JPEG encoded images is not supported"
                self._DataSyntax = 'explicit'
                self._DataEndian = 'little'
                self.PixelDataEndian = 'little'
            elif string.find(self._Value, '1.2.840.10008.1.2.5') == 0 :
                print "warning: RLE encoded images is not supported"
                self._DataSyntax = 'explicit'
                self._DataEndian = 'little'
                self.PixelDataEndian = 'little'
            else:
                # default "1.2.840.10008.1.2" and
                # "1.2.840.10008.20"
                self._DataSyntax = 'implicit'
                self._DataEndian = 'little'
                self.PixelDataEndian = 'little'

        return 1
        # print hex(self._Group), hex(self._Element), self._Length, self._DataEndian, self._MetaEndian, self._EndianChar
        
    def _DICOMReadValue(self, vr):
        """
        Load Data Element Value, unpack it, 
        """
        if (self._VR == None):
            self._VR = vr

        # some sepcial cases
        if (self._VR == 'SQ' or self._Length == 0xFFFFFFFFL): 
            return 0

        if (self._Group == 0xFFFE and not self._Encapsulated):
            return 0

        if (self._Length == 0):
            self._VM = 0
            return 0

        if (self._Length == 0xFFFFFFFFL):
            self._VM = 1
            return 0

        # assert self._Length >0, "Value length "
        # load the value
        buf = self._File.read(self._Length)


        # check file I/O error here.

        """
        for implicit VR, we have to get VR from the
        dictionary. Since dictionary has not been
        implemented yet, the following is not used
        """
        
        # unpack the value
        if ( self._VR == 'AE' or
             self._VR == 'AS' or
             self._VR == 'CS' or
             self._VR == 'DA' or
             self._VR == 'DS' or
             self._VR == 'DT' or
             self._VR == 'IS' or
             self._VR == 'LO' or
	     self._VR == 'LT' or
             self._VR == 'PN' or
             self._VR == 'SH' or
             self._VR == 'TM' or
             self._VR == 'UI' ):  # string
            self._VM = 1
            value = string.replace(buf, "\0", "" ) # remove null bytes
            # convert to a list if it contains '\'
            value = map(string.strip,string.split(value,'\\'))
            if len(value) == 1:
                value = value[0]
            self._Value = value

            del buf
            return 0

        if ( self._VR == 'AT'):  # attribute tag : pairs of (gggg,eeee)
            self._VM = self._Length / 4
            temp = struct.unpack(self._EndianChar + 2*self._VM*'H', buf)
            # group the values into pairs
            self._Value = []
            for i in range(self._VM):
                self._Value.append((temp[2*i], temp[2*i+1]))
            del buf
            return 0

        if ( self._VR == 'OW'): # word string
            self._VM = 1
            if self._DataEndian != sys_byteorder():
                tempArray = numarray.fromstring(buf,'s').byteswapped()
                buf = tempArray.tostring()
                del tempArray
            self._Value = buf
            del buf
            return 0

        value = None
        if ( self._VR == 'SS' ): # signed short
            self._VM = int(self._Length / 2)
            value = struct.unpack(self._EndianChar + self._VM*'h', buf)
            
        elif ( self._VR == 'US'): # unsigned short
            self._VM = int(self._Length / 2)
            value = struct.unpack(self._EndianChar + self._VM*'H', buf)

        elif ( self._VR == 'SL'): # singed long
            self._VM = int(self._Length / 4)
            value = struct.unpack(self._EndianChar + self._VM*'l', buf)
            
        elif ( self._VR == 'UL'): # unsigned long
            self._VM = int(self._Length / 4)
            value = struct.unpack(self._EndianChar + self._VM*'L', buf)
            
        elif ( self._VR == 'FL'): # float
            self._VM = int(self._Length / 4)
            value = struct.unpack(self._EndianChar + self._VM*'f', buf)

        elif ( self._VR == 'FD'): # double float
            self._VM = int(self._Length / 8)
            value = struct.unpack(self._EndianChar + self._VM*'d', buf)

        if value:            
            if self._VM == 1:
                self._Value = value[0]
            else:
                self._Value = value
            del value
            del buf
            return 0

        # default
        del buf
        self._VM = 1
        return 0
         
    def _DICOMSkip(self):
        if self._Length != 0xFFFFFFFFL:
            self._File.seek(self._Length, 1)

    def _ReadDICOMTags(self):
        """A replacement for _ReadDICOMInfo() that uses a DICOM dictionary.
        """
        try:
            self._DICOMOpen()
    
            self._tags = {}
            if( self._File == None ):
                return
    
            while self._DICOMReadTag():
                """
                # print tag and dictionary entry
                print "%4.4x,%4.4x" % (self._Group,self._Element), self._Length,
                try:
                    print dicomdict[(self._Group,self._Element)]
                except:
                    print 'Unrecognized'
                """
                
                entry = None
                if (self._Group & 0xff00) == 0x5000:
                    self._Group = (self._Group & 0xff00)
                elif (self._Group & 0xff00) == 0x6000:
                    self._Group = (self._Group & 0xff00)
                try:
                    entry = dicomdict[(self._Group,self._Element)]
                except KeyError:
                    pass
                    #print "Tag not found!!",(self._Group,self._Element)
                # check for pixel data
                if (self._Group > 0x7F00 and self._Group < 0x7FFF and
                    self._Group & 1 == 0 and self._Element == 0x0010):
                    self._tags['PixelDataOffset'] = self._File.tell()
                    break
                elif entry:
                    vr = entry[0]
                    if vr == 'xs':
                        vr = 'US'  # maybe 'SS' sometimes too
                    self._DICOMReadValue(vr)
                    self._tags[entry[1]] = self._Value
                else:
                    self._DICOMSkip()
    
            bits = self._tags['BitsAllocated']
            if bits == 16:
                typeCode = 's'
                numbytes = 2
            elif bits == 8:
                typeCode = '1'
                numbytes = 1
            else:
                print "No support for ", bits, " bits data yet." 
                self._PixelData = None
                numbytes = 0

            #print "number of bytes is", numbytes
            self.PixelDataType = typeCode
            self.PixelDataSwapFlag = (sys_byteorder() != self.PixelDataEndian)
            if self._LoadPixelDataFlag and numbytes:    
    
                # self._PixelData = numarray.array(typeCode)
                n = self._tags['Rows'] * self._tags['Columns']*numbytes
                print "%d data points"%n
                try:
                    n = n * self._tags['NumberOfFrames']
                    print "number of frames",self._tags['NumberOfFrames']
                except KeyError:
                    pass
                self._PixelData = numarray.fromstring( self._File.read(n), typeCode)
                if self.PixelDataSwapFlag:
                    self._PixelData.byteswap()
    
            else:
                # skip the value field
                # print "skip", self._Length
                self._File.seek(self._Length, 1)
                self._Value = None
                self._PixelData = None  
        except Exception, error:
            print "error in reading DICOM image values", error
            self._PixelData = None  
# python 2.0 has sys.byteorder
# for python 1.52, define our own
def sys_byteorder():
    if struct.pack('H',1) =='\000\001':
        return 'big'
    else:
        return 'little'
